import os
import joblib
from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score, roc_auc_score
from src.train import PowerLoadModel, data_processing
from util.commonUtil import model_view_AUC, model_hunxiao


def model_predict(test_filename):
    """
    模型预测：使用训练好的模型对测试集进行预测并评估性能
    :param test_filename: 测试数据文件名（位于../data目录下）
    """
    # 1. 初始化预测流程
    test_path = os.path.join('../data', test_filename)
    model = PowerLoadModel(test_path)
    logger = model.logfile
    logger.info(f"=======开始预测流程，测试数据：{test_path}=======")

    # 2. 加载训练好的模型和标准化器
    model_path = '../model/best_model.pkl'
    scaler_path = '../model/scaler.pkl'
    try:
        clf = joblib.load(model_path)
        scaler = joblib.load(scaler_path)
        logger.info(f"成功加载模型：{model_path} 和标准化器：{scaler_path}")
    except Exception as e:
        logger.error(f"模型或标准化器加载失败：{str(e)}", exc_info=True)
        raise

    # 3. 处理测试数据（使用训练好的标准化器）
    x_test, y_test, _ = data_processing(
        data=model.data_source,
        logger=logger,
        is_train=False,  # 测试模式
        scaler=scaler  # 传入训练好的标准化器
    )

    # 4. 模型预测
    y_pred = clf.predict(x_test)  # 预测类别
    y_pred_proba = clf.predict_proba(x_test)[:, 1]  # 预测正例概率
    logger.info("预测完成，开始计算评价指标")

    # 5. 计算评价指标
    metrics = {
        '准确率': accuracy_score(y_test, y_pred),
        '精确率': precision_score(y_test, y_pred),
        '召回率': recall_score(y_test, y_pred),
        'F1分数': f1_score(y_test, y_pred),
        'AUC': roc_auc_score(y_test, y_pred_proba)
    }

    # 打印并记录指标
    metrics_str = "\n".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
    print(f"评价指标：\n{metrics_str}")
    logger.info(f"评价指标：\n{metrics_str}")

    # 6. 可视化评估结果
    logger.info("绘制ROC曲线和混淆矩阵...")
    model_view_AUC(x_test, y_test, clf)  # ROC曲线
    model_hunxiao(y_test, y_pred)  # 混淆矩阵

    logger.info("=======预测流程结束=======")


if __name__ == '__main__':
    # 单独运行预测时使用
    model_predict('test.csv')